import cv2
import numpy as np
from myUtils.utils import xyxy_to_xywh


class KalmanTrack:
    def __init__(self):
        # 状态数和观测数需要修改
        self.kalman = cv2.KalmanFilter(6, 4)  # 6：状态数，包括（xmin，ymin，xmax，ymax, dx，dy）坐标及速度（每次移动的距离）；
                                              # 4：观测量，能看到的是坐标值

        # 状态转移矩阵
        A = np.array([[1, 0, 0, 0, 1, 0],
                      [0, 1, 0, 0, 0, 1],
                      [0, 0, 1, 0, 0, 0],
                      [0, 0, 0, 1, 0, 0],
                      [0, 0, 0, 0, 1, 0],
                      [0, 0, 0, 0, 0, 1]], np.float32)
        self.kalman.transitionMatrix = A  # 状态转移矩阵

        # 控制矩阵
        B = None
        self.kalman.controlMatrix = B

        # 状态观测矩阵
        H = np.array([[1, 0, 0, 0, 0, 0],
                      [0, 1, 0, 0, 0, 0],
                      [0, 0, 1, 0, 0, 0],
                      [0, 0, 0, 1, 0, 0]], np.float32)
        self.kalman.measurementMatrix = H  # 系统测量矩阵

        # 观测噪声协方差矩阵R，p(v)~N(0,R)
        # 观测噪声来自于检测框丢失、重叠等
        R = np.array([[1, 0, 0, 0],
                      [0, 1, 0, 0],
                      [0, 0, 1, 0],
                      [0, 0, 0, 1]], np.float32) * 0.001  # 越小
        self.kalman.measurementNoiseCov = R

        # 过程噪声协方差矩阵Q，p(w)~N(0,Q)，噪声来自真实世界中的不确定性,
        # 在跟踪任务当中，过程噪声来自于目标移动的不确定性（突然加速、减速、转弯等）
        Q = np.eye(6, dtype=np.float32) * 0.001
        self.kalman.processNoiseCov = Q

        # 状态估计协方差矩阵P初始化
        P = np.eye(6, dtype=np.float32)
        self.kalman.errorCovPre = P

        self.cur_measurement = np.nan
        self.cur_prediction = np.nan

        self.pre_measurement = np.nan
        self.pre_prediction = np.nan

    def get_cur_state(self, target_box):
        """获取初始值状态测量值"""
        # target_box = [729, 238, 764, 339]
        self.cur_measurement = target_box  # 目标初始bouding box

        self.cur_measurement = xyxy_to_xywh(self.cur_measurement)

        #  [中心x,中心y,宽w,高h]
        self.cur_measurement = np.array(
            [[np.float32(self.cur_measurement[0]), np.float32(self.cur_measurement[1]),
              np.float32(self.cur_measurement[2]), np.float32(self.cur_measurement[3])]]).T

        return self.cur_measurement

    def get_initial_state(self, target_box):
        self.cur_measurement = self.get_cur_state(target_box)

        self.pre_measurement = self.cur_measurement
        self.cur_prediction = self.cur_measurement
        self.pre_prediction = self.cur_measurement

    def correct_and_predict(self, target_box, iou_matched):
        # 将前线状态进行存储
        self.pre_measurement = self.cur_measurement
        self.pre_prediction = self.cur_prediction

        # 用当前测量来校正卡尔曼滤波器

        self.cur_measurement = self.pre_prediction[0: 4, :]
        if iou_matched:  # 根据IOU来计算
            self.cur_measurement = self.get_cur_state(target_box)

        self.kalman.correct(self.cur_measurement)  # 用当前测量来校正卡尔曼滤波器
        self.cur_prediction = self.kalman.predict()  # 计算卡尔曼预测值，作为当前预测
        return self.cur_prediction


if __name__ == '__main__':
    kalman_tracker = KalmanTrack()
    kalman_tracker.get_initial_state([729, 288, 35, 101])  # xywh

    while True:
        # 将先前的预测值作为当前的测量值
        data = list([kalman_tracker.pre_prediction[0][0], kalman_tracker.pre_prediction[1][0],
                    kalman_tracker.pre_prediction[2][0], kalman_tracker.pre_prediction[3][0]])
        print(kalman_tracker.correct_and_predict(data))









